#include <cuda.h>
#include <iostream>

#include "randomc.h"
#include "stocc.h"

#ifndef MULTIFILE_PROJECT
// If compiled as a single file then include these cpp files, 
// If compiled as a project then compile and link in these cpp files.
   #include "mersenne.cpp"             // code for random number generator
   #include "stoc1.cpp"                // random library source code
   #include "userintf.cpp"             // define system specific user interface
#endif

#define BLOCKSIZE_X 1
#define BLOCKSIZE_Y 512

#define BLOCKSIZE_Tmax 512

__constant__ float Y_d[1000];


void Yperm(float *Y, const int N)
{
  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  for (int n = 1; n < N; n++){
    int idx = sto.IRandom(0, n); // uniform over [0,n]
    float val = Y[n];
    Y[n] = Y[idx];
    Y[idx] = val;
  }
}

// prototypes
__global__ void MeanX(const float *X, float *meanX, const int N, const int K, const size_t pitch_x)
{
  __shared__ float mean[BLOCKSIZE_Y][BLOCKSIZE_X];

  int idx = threadIdx.x;
  int idy = threadIdx.y;

  int colId = blockIdx.x*BLOCKSIZE_Y + idy;

  if (colId < K){
    mean[idy][idx] = 0.0;

    // compute mean
    for (int n = idx; n < N; n += BLOCKSIZE_X)   
      mean[idy][idx] += *((float*)((char*)X + n * pitch_x) + colId);
    __syncthreads();

    // add partial quantities
    if (BLOCKSIZE_X >= 512){
      if (idx < 256)
	mean[idy][idx] += mean[idy][idx + 256];
      __syncthreads();
    }

    if (BLOCKSIZE_X >= 256){
      if (idx < 128)
	mean[idy][idx] += mean[idy][idx + 128];
      __syncthreads();
    }

    if (BLOCKSIZE_X >= 128){
      if (idx < 64)
	mean[idy][idx] += mean[idy][idx + 64];
      __syncthreads();
    }

    if (BLOCKSIZE_X >= 64){
      if (idx < 32)
	mean[idy][idx] += mean[idy][idx + 32];
    }

    if (idx < 16){
      if (BLOCKSIZE_X >= 32)
	mean[idy][idx] += mean[idy][idx + 16];

      if (BLOCKSIZE_X >= 16)
	mean[idy][idx] += mean[idy][idx +  8];

      if (BLOCKSIZE_X >= 8)
	mean[idy][idx] += mean[idy][idx +  4];

      if (BLOCKSIZE_X >= 4)
	mean[idy][idx] += mean[idy][idx +  2];

      if (BLOCKSIZE_X >= 2)
	mean[idy][idx] += mean[idy][idx +  1];
    }

    if (idx == 0)
      meanX[colId] = mean[idy][0] / N;
  }
}

__global__ void SumQuadX(const float *X, float *meanX, float *sumquadX, const int N, const int K, const size_t pitch_x)
{
  __shared__ float aux[BLOCKSIZE_Y][BLOCKSIZE_X];

  int idx = threadIdx.x;
  int idy = threadIdx.y;

  int colId = blockIdx.x*BLOCKSIZE_Y + idy;

  float xn;
  float xk;

  if (colId < K){
    aux[idy][idx] = 0.0;

    // compute mean
    xk = meanX[colId];
    for (int n = idx; n < N; n += BLOCKSIZE_X){   
      xn = *((float*)((char*)X + n * pitch_x) + colId);
      aux[idy][idx] += (xn-xk) * (xn-xk);
    }

    __syncthreads();

    // add partial quantities
    if (BLOCKSIZE_X >= 512){
      if (idx < 256)
	aux[idy][idx] += aux[idy][idx + 256];
      __syncthreads();
    }

    if (BLOCKSIZE_X >= 256){
      if (idx < 128)
	aux[idy][idx] += aux[idy][idx + 128];
      __syncthreads();
    }

    if (BLOCKSIZE_X >= 128){
      if (idx < 64)
	aux[idy][idx] += aux[idy][idx + 64];
      __syncthreads();
    }

    if (BLOCKSIZE_X >= 64){
      if (idx < 32)
	aux[idy][idx] += aux[idy][idx + 32];
    }

    if (idx < 16){
      if (BLOCKSIZE_X >= 32)
	aux[idy][idx] += aux[idy][idx + 16];

      if (BLOCKSIZE_X >= 16)
	aux[idy][idx] += aux[idy][idx +  8];

      if (BLOCKSIZE_X >= 8)
	aux[idy][idx] += aux[idy][idx +  4];

      if (BLOCKSIZE_X >= 4)
	aux[idy][idx] += aux[idy][idx +  2];

      if (BLOCKSIZE_X >= 2)
	aux[idy][idx] += aux[idy][idx +  1];
    }

    if (idx == 0)
      sumquadX[colId] = aux[idy][0];
  }
}

__global__ void Tstat(const float *X, float *T, const float *meanX, const float *sumquadX, const int N, const int K, const float meanY, const float sumquadY, const size_t pitch_x)
{
  int colId = blockIdx.x*blockDim.x + threadIdx.x;
  
  if(colId < K){
    float stat = 0.0;
    // compute the test statistic 
    for (int n = 0; n < N; n ++)    
      stat += Y_d[n] * (*((float*)((char*)X + n * pitch_x) + colId));

    stat = (stat-N*meanY*meanX[colId])/sqrt(sumquadX[colId]*sumquadY);

    // store result to global variable
    T[colId] = stat*sqrt(N-2.0)/sqrt(1.0-stat*stat);
  }
}

__global__ void Tmax(const float *Ttilde, const int K, float *Tmax)
{
  __shared__ float zeta [BLOCKSIZE_Tmax];

  int idx = threadIdx.x;
  zeta[idx] = 0.0;
  for (int k = idx; k < K; k += BLOCKSIZE_Tmax){
    if (zeta[idx] < abs(Ttilde[k]))
      zeta[idx] = abs(Ttilde[k]);
  }

  __syncthreads();

  // compute the global maximum
  if (BLOCKSIZE_Tmax >= 512){
    if (idx < 256){
      if (zeta[idx] < zeta[idx+256])
	zeta[idx] = zeta[idx+256];
    }
    __syncthreads();
  }

  if (BLOCKSIZE_Tmax >= 256){
    if (idx < 128){
      if (zeta[idx] < zeta[idx+128])
	zeta[idx] = zeta[idx+128];
    }
    __syncthreads();
  }

  if (BLOCKSIZE_Tmax >= 128){
    if (idx < 64){
      if (zeta[idx] < zeta[idx+64])
	zeta[idx] = zeta[idx+64];
    }
    __syncthreads();
  }

  if (BLOCKSIZE_Tmax >= 64){
    if (idx < 32){
      if (zeta[idx] < zeta[idx+32])
	zeta[idx] = zeta[idx+32];
    }
  }

  if (idx < 16){
    if (BLOCKSIZE_Tmax >= 32){
      if (zeta[idx] < zeta[idx+16])
	zeta[idx] = zeta[idx+16];
    }
 
    if (BLOCKSIZE_Tmax >= 16){
      if (zeta[idx] < zeta[idx+8])
	zeta[idx] = zeta[idx+8];
    }
  
    if (BLOCKSIZE_Tmax >= 8){
      if (zeta[idx] < zeta[idx+4])
	zeta[idx] = zeta[idx+4];
    }
 
    if (BLOCKSIZE_Tmax >= 4){
      if (zeta[idx] < zeta[idx+2])
	zeta[idx] = zeta[idx+2];
    }
   
    if (BLOCKSIZE_Tmax >= 2){
      if (zeta[idx] < zeta[idx+1])
	zeta[idx] = zeta[idx+1];
    }
  }

  // store result to global variable
  if (idx == 0)
    Tmax[0] = zeta[0];
}

__global__ void Pvalue(const float *T, const float *Ttilde, float *p, float *P, const float *Tmax, const int K){
  int idx = blockIdx.x*blockDim.x + threadIdx.x;
 
  if(idx < K){
    if(abs(Ttilde[idx]) >= abs(T[idx]))
      p[idx]++;
    if(Tmax[0] >= abs(T[idx]))
      P[idx]++;
  }
}

long int seed = (long int)time(0); // global seed

int main(int argc, const char *argv[])
{
  // argv[1]: N
  // argv[2]: K
  // argv[3]: B

  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);

  int N = atoi(argv[1]);
  int K = atoi(argv[2]);
  int B = atoi(argv[3]);

  // print
  std::cout<<"(N="<<N<<",K="<<K<<",B="<<B<<")"<<"\n";

  // allocate memory on host  
  float *X_h;                                                                                                                         
  X_h = (float *)malloc(N*K*sizeof(float));
  float *Y_h;
  Y_h = (float *)malloc(N*sizeof(float));
  float *p_h;
  p_h = (float *)malloc(K*sizeof(float));
  float *P_h;
  P_h = (float *)malloc(K*sizeof(float));
  float *T_h;
  T_h = (float *)malloc(K*sizeof(float));
  float *Ttilde_h;
  Ttilde_h = (float *)malloc(K*sizeof(float));

  int count = 0;
  for (int n = 0; n < N; n++){
    Y_h[n] = sto.Normal(0, 1.0);
    for (int k = 0; k < K; k++){
      X_h[count] = sto.Normal(0, 1.0);
      count++;
    }
  }

  // block until the kernels complete
  cudaError_t error = cudaThreadSynchronize();
  if (error != cudaSuccess){
    std::cout << "Failed at ThreadSync:TimeBegin\n";
    return 0;
  }

  // start time profiling                                                                                                                 
  cudaEvent_t start_cuda, stop;
  float elapsed;
  cudaThreadSynchronize();
  cudaEventCreate(&start_cuda);
  cudaEventCreate(&stop);
  cudaEventRecord(start_cuda, 0);

  for (int k = 0; k < K; k++){
    p_h[k] = 0;
    P_h[k] = 0;
    T_h[k] = 0.0;
    Ttilde_h[k] = 0.0;
  }

  size_t pitch_x;

  // allocate memory on device    
  float *X_d;
  error = cudaMallocPitch((void **) &X_d, &pitch_x, K*sizeof(float), N);
  if (error != cudaSuccess){
    std::cout << "Failed at malloc:X_d.\n";
    return 0;
  }
  error = cudaMemcpy2D(X_d, pitch_x, X_h, K*sizeof(float), K*sizeof(float), N, cudaMemcpyHostToDevice);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:X_d.\n";
    return 0;
  }

  error = cudaMemcpyToSymbol(Y_d, Y_h, N*sizeof(int), 0, cudaMemcpyHostToDevice);
  if (error != cudaSuccess){
    std::cout << "Failed at copy to symbol:Y_d.\n";
    return 0;
  }

  float *meanX_d;
  error = cudaMalloc((void **) &meanX_d, K*sizeof(float));
  if (error != cudaSuccess){
    std::cout<< "Failed at malloc:meanX_d.\n";
    return 0;
  }

  float *sumquadX_d;
  error = cudaMalloc((void **) &sumquadX_d, K*sizeof(float));
  if (error != cudaSuccess){
    std::cout<< "Failed at malloc:sumquadX_d.\n";
    return 0;
  }

  float *p_d;
  error = cudaMalloc((void **) &p_d, K*sizeof(float));
  if (error != cudaSuccess){
    std::cout << "Failed at malloc:p_d.\n";
    return 0;
  }
  error = cudaMemcpy(p_d, p_h, sizeof(float)*K, cudaMemcpyHostToDevice);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:p_d.\n";
    return 0;
  }

  float *P_d;
  error = cudaMalloc((void **) &P_d, K*sizeof(float));
  if (error != cudaSuccess){
    std::cout << "Failed at malloc:P_d.\n";
    return 0;
  }
  error = cudaMemcpy(P_d, P_h, sizeof(float)*K, cudaMemcpyHostToDevice);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:P_d.\n";
    return 0;
  }

  float *T_d;
  error = cudaMalloc((void **) &T_d, K*sizeof(float));
  if (error != cudaSuccess){
    std::cout << "Failed at malloc:T_d.\n";
    return 0;
  }
  error = cudaMemcpy(T_d, T_h, sizeof(float)*K, cudaMemcpyHostToDevice);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:T_d.\n";
    return 0;
  }

  float *Ttilde_d;
  error = cudaMalloc((void **) &Ttilde_d, K*sizeof(float));
  if (error != cudaSuccess){
    std::cout << "Failed at malloc:Ttilde_d.\n";
    return 0;
  }
  error = cudaMemcpy(Ttilde_d, Ttilde_h, sizeof(float)*K, cudaMemcpyHostToDevice);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:Ttilde_d.\n";
    return 0;
  }

  float *Tmax_d;
  error = cudaMalloc((void **) &Tmax_d, sizeof(float));
  if (error != cudaSuccess){
    std::cout << "Failed at malloc:Tmax_d.\n";
    return 0;
  }
  
  // configuration
  int blockSize = 512;
  int nBlocks = K/blockSize + (K%blockSize == 0?0:1);

  dim3 threads(BLOCKSIZE_X, BLOCKSIZE_Y, 1);
  dim3 grid(K/BLOCKSIZE_Y + (K%BLOCKSIZE_Y == 0?0:1), 1, 1);

  // compute mean of Y
  float meanY = 0.0;
  for (int n = 0; n < N; n++)
    meanY += Y_h[n];
  meanY /= N;

  // compute sumquadY
  float sumquadY = 0.0;
  for (int n = 0; n < N; n++)
    sumquadY += (Y_h[n]-meanY) * (Y_h[n]-meanY);

  // compute meanX
  MeanX<<<grid, threads>>>(X_d, meanX_d, N, K, pitch_x);

  // compute sumquadX
  SumQuadX<<<grid, threads>>>(X_d, meanX_d, sumquadX_d, N, K, pitch_x);

  // compute T
  Tstat<<<nBlocks, blockSize>>>(X_d, T_d, meanX_d, sumquadX_d, N, K, meanY, sumquadY, pitch_x);

  for (int b = 0; b < B; b++){
    // permute Y
    Yperm(Y_h, N);
    error = cudaMemcpyToSymbol(Y_d, Y_h, N*sizeof(int), 0, cudaMemcpyHostToDevice);
    if (error != cudaSuccess){
      std::cout << "Failed at copy to symbol:Y_d.\n";
      return 0;
    }

    // compute test statistic
    Tstat<<<nBlocks, blockSize>>>(X_d, Ttilde_d, meanX_d, sumquadX_d, N, K, meanY, sumquadY, pitch_x);

    // compute maximum of test statistic
    Tmax <<<1, BLOCKSIZE_Tmax>>> (Ttilde_d, K, Tmax_d);

    // compute p-values
    Pvalue <<<nBlocks, blockSize>>> (T_d, Ttilde_d, p_d, P_d, Tmax_d, K);
  }

  // retrieve result from device                                                                                                                                          
  error = cudaMemcpy(p_h, p_d, sizeof(float)*K, cudaMemcpyDeviceToHost);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:p_h.\n";
    return 0;
  }

  error = cudaMemcpy(P_h, P_d, sizeof(float)*K, cudaMemcpyDeviceToHost);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:P_h.\n";
    return 0;
  }

  error = cudaMemcpy(T_h, T_d, sizeof(float)*K, cudaMemcpyDeviceToHost);
  if (error != cudaSuccess){
    std::cout << "Failed at memcpy:T_h.\n";
    return 0;
  }

  // block until the kernels complete
  error = cudaThreadSynchronize();
  if (error != cudaSuccess){
    std::cout << "Failed at ThreadSync:TimeEnd\n";
    return 0;
  }

  cudaThreadSynchronize();
  cudaEventRecord(stop, 0);
  cudaEventSynchronize(stop);
  cudaEventElapsedTime(&elapsed, start_cuda, stop);

  cudaEventDestroy(start_cuda);
  cudaEventDestroy(stop);

  std::cout<<"elapsed time:"<<elapsed<<" milliseconds"<<"\n";

  // cleanup
  free(Y_h);
  free(X_h);
  free(p_h);
  free(P_h);
  free(T_h);
  free(Ttilde_h);

  error = cudaFree(X_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:X_d.\n";
    return 0;
  }
     
  error = cudaFree(p_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:p_d.\n";
    return 0;
  }
      
  error = cudaFree(P_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:P_d.\n";
    return 0;
  }
     
  error = cudaFree(T_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:T_d.\n";
    return 0;
  }
     
  error = cudaFree(Ttilde_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:Ttilde_d.\n";
    return 0;
  }
    
  error = cudaFree(Tmax_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:Tmax_d.\n";
    return 0;
  }

  error = cudaFree(meanX_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:meanX_d.\n";
    return 0;
  }

  error = cudaFree(sumquadX_d);
  if (error != cudaSuccess){
    std::cout << "Failed at free:sumquadX_d.\n";
    return 0;
  }

  return(0);
}
